# -*- ecoding: utf-8 -*-
# @ModuleName: windows
# @Function: 
# @Author: Liweijian
# @Time: 2024/12/20 15:26

import shutil
import sys

import cv2
import numpy as np
import tensorflow as tf
from PIL import Image
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *


class MainWindow(QTabWidget):
    # 初始化
    def __init__(self):
        super().__init__()
        self.setWindowIcon(QIcon('images/APP_Design_Pictures/logo.jpg'))
        self.setWindowTitle('基于深度学习的图像分类系统DesignBy李唯戬')  # 系统名称
        # 模型初始化
        self.model = tf.keras.models.load_model("D:/pythonProject/Final_MNIST/models/cnn_mnist.h5")
        self.to_predict_name = "images/APP_Design_Pictures/img.png"
        self.class_names = [str(i) for i in range(10)]  # MNIST 类别名称（0-9）
        self.resize(1000, 900)
        self.timer_camera = QTimer()  # 创建一个定时器
        self.video_capture = cv2.VideoCapture()  # 创建一个视频捕获对象
        self.CAM_NUM = 0  # 摄像头编号，默认为 0
        self.initUI()

    # 界面初始化，设置界面布局
    def initUI(self):
        # 主页面，整体布局设置
        main_widget = QWidget()
        main_layout = QHBoxLayout()  # 水平布局
        font = QFont('黑体', 15)

        # 左侧布局
        left_widget = QWidget()
        left_layout = QVBoxLayout()  # 垂直布局
        img_title = QLabel("图片显示区")
        img_title.setFont(font)
        img_title.setAlignment(Qt.AlignCenter)
        video_title = QLabel("视频显示区")
        video_title.setFont(font)
        video_title.setAlignment(Qt.AlignCenter)

        self.img_label = QLabel()
        img_init = cv2.imread(self.to_predict_name)
        h, w, c = img_init.shape
        scale = 400 / h
        img_show = cv2.resize(img_init, (0, 0), fx=scale, fy=scale)
        cv2.imwrite("images/show.png", img_show)
        img_init = cv2.resize(img_init, (28, 28))  # 改为28x28大小
        cv2.imwrite('images/target.png', img_init)
        self.img_label.setPixmap(QPixmap("images/show.png"))

        self.label_video = QLabel(self)
        self.label_video.setGeometry(40, 480, 400*scale, 400*scale)  # 设置 QLabel 的位置和大小
        self.label_video.setStyleSheet("border: 1px solid gray")

        # 添加组件到左侧布局
        left_layout.addWidget(img_title)
        left_layout.addWidget(self.img_label, Qt.AlignTop)
        left_layout.addWidget(video_title)
        left_layout.addWidget(self.label_video, Qt.AlignBottom)
        left_widget.setLayout(left_layout)

        # 右侧布局
        right_widget = QWidget()
        right_layout = QVBoxLayout()

        self.btn_open = QPushButton(" 打开摄像头 ")
        self.btn_open.clicked.connect(self.display_video)
        self.btn_open.setFont(font)

        self.btn_capture = QPushButton('拍照')
        self.btn_capture.clicked.connect(self.capture_image)
        self.btn_capture.setFont(font)

        btn_change = QPushButton(" 上传图片 ")
        btn_change.clicked.connect(self.change_img)
        btn_change.setFont(font)

        btn_predict = QPushButton(" 开始识别 ")
        btn_predict.setFont(font)
        btn_predict.clicked.connect(self.predict_img)

        label_result = QLabel(' 识别的数字： ')
        self.result = QLabel("等待识别")
        label_result.setFont(QFont('黑体', 16))
        self.result.setFont(QFont('黑体', 24))

        # 添加组件到右侧布局
        right_layout.addStretch()
        right_layout.addWidget(label_result, 0, Qt.AlignCenter)
        right_layout.addStretch()
        right_layout.addWidget(self.result, 0, Qt.AlignCenter)
        right_layout.addStretch()
        right_layout.addWidget(self.btn_open)
        right_layout.addWidget(self.btn_capture)
        right_layout.addWidget(btn_change)
        right_layout.addWidget(btn_predict)
        right_layout.addStretch()

        right_widget.setLayout(right_layout)

        # 主布局
        main_layout.addWidget(left_widget)
        main_layout.addWidget(right_widget)
        main_widget.setLayout(main_layout)

        self.timer_camera.timeout.connect(self.show_camera)

        # 关于页面
        about_widget = QWidget()
        about_layout = QVBoxLayout()
        about_title = QLabel('基于深度学习的图像分类系统')
        about_title.setFont(QFont('黑体', 18))
        about_title.setAlignment(Qt.AlignCenter)
        about_img = QLabel()
        about_img.setPixmap(QPixmap('images/APP_Design_Pictures/logo.jpg'))
        about_img.setAlignment(Qt.AlignCenter)
        label_super = QLabel("24-25学年python期末大项目 @202234070359李唯戬")
        label_super.setFont(QFont('黑体', 12))
        label_super.setAlignment(Qt.AlignRight)

        about_layout.addWidget(about_title)
        about_layout.addStretch()
        about_layout.addWidget(about_img)
        about_layout.addStretch()
        about_layout.addWidget(label_super)
        about_widget.setLayout(about_layout)

        self.addTab(main_widget, '主页')
        self.addTab(about_widget, '关于')
        self.setTabIcon(0, QIcon('images/APP_Design_Pictures/主页面.png'))
        self.setTabIcon(1, QIcon('images/APP_Design_Pictures/关于.png'))

    # 开启/关闭摄像头
    def display_video(self):
        if not self.timer_camera.isActive():
            flag = self.video_capture.open(self.CAM_NUM)
            print(flag)
            if not flag:
                QMessageBox.warning(self, '警告！', "无法开启摄像头，请检查！")
            else:
                print("摄像头成功开启。")
                self.timer_camera.start(30)  # 启动定时器
                self.btn_open.setText('关闭摄像头')
        else:
            self.timer_camera.stop()
            self.video_capture.release()
            self.label_video.clear()
            self.btn_open.setText('打开摄像头')

    # 显示视频
    def show_camera(self):
        ret, frame = self.video_capture.read()
        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            image = QImage(frame.data, frame.shape[1], frame.shape[0], QImage.Format_RGB888)
            self.label_video.setPixmap(QPixmap.fromImage(image))

    # 获取摄像头捕捉的图片
    def capture_image(self):
        if self.video_capture.isOpened():
            ret, frame = self.video_capture.read()
            if ret:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                image = QImage(frame.data, frame.shape[1], frame.shape[0], QImage.Format_RGB888)
                self.img_label.setPixmap(QPixmap.fromImage(image))
                # 将图像调整为28x28
                img_init = cv2.resize(frame, (28, 28))
                cv2.imwrite('images/target.png', img_init)  # 保存目标图片
                self.result.setText("等待识别")

    # 上传并显示图片
    def change_img(self):
        openfile_name = QFileDialog.getOpenFileName(self, '选择文件', '', 'Image files(*.jpg *.png *jpeg)')
        img_name = openfile_name[0]
        if img_name:
            target_image_name = "images/tmp_up." + img_name.split(".")[-1]
            shutil.copy(img_name, target_image_name)
            self.to_predict_name = target_image_name
            img_init = cv2.imread(self.to_predict_name)
            img_init = cv2.cvtColor(img_init, cv2.COLOR_BGR2GRAY)  # 转换为灰度图
            img_init = cv2.resize(img_init, (28, 28))  # 调整到28x28大小
            cv2.imwrite('images/target.png', img_init)  # 保存目标图片
            img_show = cv2.resize(img_init, (0, 0), fx=10, fy=10)  # 显示用的缩略图
            cv2.imwrite("images/show.png", img_show)
            self.img_label.setPixmap(QPixmap("images/show.png"))
            self.result.setText("等待识别")

    # 预测图片
    def predict_img(self):
        img = Image.open('images/target.png')  # 读取图片
        img = img.convert('L')  # 转换为灰度图
        img = img.resize((28, 28))  # 调整大小为28x28
        img_array = np.asarray(img)  # 转换为numpy数组
        img_array = img_array / 255.0  # 归一化到[0, 1]
        img_array = img_array.reshape(1, 28, 28, 1)  # 调整形状以符合模型输入

        outputs = self.model.predict(img_array)  # 进行预测
        result_index = int(np.argmax(outputs))  # 获取结果索引
        result = self.class_names[result_index]  # 获取对应的类别名称
        self.result.setText(result)  # 在界面上显示结果

    # 界面关闭事件
    def closeEvent(self, event):
        reply = QMessageBox.question(self, '退出', "是否要退出程序？", QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        if reply == QMessageBox.Yes:
            self.close()
            event.accept()
        else:
            event.ignore()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    x = MainWindow()
    x.show()
    sys.exit(app.exec_())
